# Read in the timevar file produced by snoopy and analyze it.

from numpy import *

def read(tvfile):
    global cols, data
    
    try:
        fid = open(tvfile,'r')
    except IOError as e:
        print "I/O error({0}): {1}".format(e.errno, e.strerror)

    hdrs = fid.readline().split()
    cols = { hdrs[k] : k for k in range(len(hdrs))}
    data = loadtxt(fid,skiprows=1)
    print 'The data in '+tvfile+' is in {:d} rows by {:d} columns'.format(*shape(data))
    fid.close()

    return cols,data

def firstrow(tmin=None):
    if tmin == None :
        i = 0
    else :
        t = data[:,cols['t']]
        i = abs(tmin-t).argmin()
    return i

def ev_avg(tmin=None):
    im = firstrow(tmin)
    try:
       ev = data[:,cols['ev']]
       avg = ev[im:].mean()
       return avg
    except KeyError :
        print 'Could not find a column label "ev"'

def em_avg(tmin=None):
    im = firstrow(tmin)
    try:
       em = data[:,cols['em']]
       avg = em[im:].mean()
       return avg
    except KeyError :
        print 'Could not find a column label "em"'

def vstress(tmin=None):
    im = firstrow(tmin)
    vtotsq = 2*data[im:,cols['ev']].mean()
    vxsq   = data[im:,cols['vxvx']].mean()
    vzsq   = data[im:,cols['vzvz']].mean()
    vxvy   = data[im:,cols['vxvy']].mean()    
    vysq = vtotsq - vxsq - vzsq
    return array([[vxsq,vxvy,0.],[vxvy,vysq,0.],[0.,0.,vzsq]])

def bstress(tmin=None):
    im = firstrow(tmin)
    btotsq = 2*data[im:,cols['em']].mean()
    bxsq   = data[im:,cols['bxbx']].mean()
    bzsq   = data[im:,cols['bzbz']].mean()
    bxby   = data[im:,cols['bxby']].mean()    
    bysq = btotsq - bxsq - bzsq
    return array([[bxsq,bxby,0.],[bxby,bysq,0.],[0.,0.,bzsq]])

